import pandas as pd
# import datetime
# import pymongo
import tushare as ts
# import matplotlib as mpl
import matplotlib.pyplot as plt


def calculate_beta(df_stock,
                   df_hs300,
                   code,
                   start='2017-01-01',
                   end='2017-11-20'):
    df = pd.DataFrame(
        {
            code: df_stock['close'].pct_change(),
            'hs300': df_hs300['close'].pct_change()
        },
        index=df_stock.index)
    cov = df.corr().iloc[0, 1]
    df_hs300['change'] = df_hs300['close'].pct_change() * 100
    var = df_hs300['change'].var()
    beta = cov / var
    return beta


# 定价曲线
def make_capm(code, start='2019-01-01', end='2019-11-20'):
    df_stock = ts.get_hist_data(code, start=start, end=end)
    df_hs300 = ts.get_hist_data('hs300', start=start, end=end)
    df_stock.sort_index(ascending=True, inplace=True)
    df_hs300.sort_index(ascending=True, inplace=True)
    beta = calculate_beta(df_stock, df_hs300, code, start=start, end=end)
    loss_free_return = 0.04
    df = pd.DataFrame(
        {
            code: df_stock['close'] / df_stock['close'].values[1] - 1,
            'hs300': df_hs300['close'] / df_hs300['close'].values[1] - 1,
            'days': range(1, df_stock.shape[0] + 1)
        },
        index=df_stock.index)
    # 一年大致含有大概250个交易日，所以这里是求出年化收益平均到每个交易日上的收益
    df['beta'] = df['days'] * loss_free_return / 250 + beta * (
        df['hs300'] - df['days'] * loss_free_return / 250)
    df['alpha'] = df[code] - df['beta']
    df[[code, 'hs300', 'beta', 'alpha']].plot(figsize=(960 / 72, 480 / 72))
    plt.show()


# make_capm('601318')

# make_capm('600030')
# make_capm('600030',start='2017-10-10')
# make_capm('600036')
make_capm('601688')
